package com.shujia.ml

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo6FitImageModel {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .appName("point")
      .getOrCreate()
    import spark.implicits._

    //1、读取图片的数据
    val imageData: DataFrame = spark
      .read
      .format("libsvm")
      .option("numFeatures", 784)
      .load("spark/data/image_data")

    //2、将数据拆分成训练集和测试集
    val Array(train, test) = imageData.randomSplit(Array(0.8, 0.2))

    //3、选择算法
    val regression = new LogisticRegression()

    //4、将训练集带入算法训练模型
    val model: LogisticRegressionModel = regression.fit(train)

    //5、将测试集带入模型测试模型准确率
    val frame: DataFrame = model.transform(test)

    //6、计算准确率
    val p: Double = frame.where($"label" === $"prediction").count().toDouble / frame.count()
    println(p)

    //7、保存模型
    model.save("spark/data/image_model")
  }
}
